"""
常用代码模板库 - 为常见场景提供高质量模板
用于在生成失败时提供回退方案
"""
from typing import Optional

TEMPLATES = {
    "rotating_file_logger": '''import logging
from logging.handlers import RotatingFileHandler
from typing import Optional

def setup_rotating_logger(
    log_file: str,
    max_bytes: int = 10485760,
    backup_count: int = 5,
    level: int = logging.ERROR
) -> logging.Logger:
    """
    Configure a logger with rotating file handler.
    
    Args:
        log_file: Path to log file
        max_bytes: Max file size in bytes before rotation (default: 10MB)
        backup_count: Number of backup files to keep
        level: Logging level (default: ERROR)
    
    Returns:
        Configured logger instance
    """
    logger = logging.getLogger(__name__)
    logger.setLevel(level)
    
    handler = RotatingFileHandler(
        log_file,
        maxBytes=max_bytes,
        backupCount=backup_count
    )
    
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    
    return logger
''',

    "http_retry_request": '''import time
import requests
from typing import Dict, Any, Optional

def fetch_with_retry(
    url: str,
    retries: int = 3,
    timeout: float = 5.0,
    backoff: float = 0.5
) -> Dict[str, Any]:
    """
    Fetch JSON from URL with retry and exponential backoff.
    
    Args:
        url: Target URL
        retries: Number of retry attempts
        timeout: Request timeout in seconds
        backoff: Initial backoff delay in seconds
    
    Returns:
        Parsed JSON response as dict
    
    Raises:
        requests.HTTPError: If all retries fail
    """
    last_error = None
    
    for attempt in range(1, retries + 1):
        try:
            response = requests.get(url, timeout=timeout)
            response.raise_for_status()
            return response.json()
        except (requests.RequestException, ValueError) as e:
            last_error = e
            if attempt < retries:
                delay = backoff * (2 ** (attempt - 1))
                time.sleep(delay)
    
    raise last_error
''',

    "pandas_filter_csv": '''import pandas as pd
from typing import Any, Optional

def filter_csv_rows(
    file_path: str,
    column: str,
    target_value: Any,
    top_n: int = 10
) -> pd.DataFrame:
    """
    Read CSV and return top N rows where column equals target value.
    
    Args:
        file_path: Path to CSV file
        column: Column name to filter on
        target_value: Value to match
        top_n: Maximum number of rows to return
    
    Returns:
        Filtered DataFrame with top N matching rows
    """
    df = pd.read_csv(file_path)
    filtered = df[df[column] == target_value]
    return filtered.head(top_n)
''',

    "sqlite_database": '''import sqlite3
from typing import List, Dict, Any, Optional
from contextlib import contextmanager

@contextmanager
def get_db_connection(db_path: str):
    """
    Context manager for SQLite database connection.
    
    Args:
        db_path: Path to SQLite database file
    
    Yields:
        Database connection object
    """
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

def execute_query(
    db_path: str,
    query: str,
    params: Optional[tuple] = None
) -> List[Dict[str, Any]]:
    """
    Execute SQL query and return results as list of dicts.
    
    Args:
        db_path: Path to database file
        query: SQL query string
        params: Query parameters tuple
    
    Returns:
        List of row dictionaries
    """
    with get_db_connection(db_path) as conn:
        cursor = conn.cursor()
        cursor.execute(query, params or ())
        return [dict(row) for row in cursor.fetchall()]
''',

    "file_encryption": '''from cryptography.fernet import Fernet
from typing import Optional
import os

def encrypt_file(
    input_path: str,
    output_path: str,
    key: Optional[bytes] = None
) -> bytes:
    """
    Encrypt file using Fernet symmetric encryption.
    
    Args:
        input_path: Path to input file
        output_path: Path to save encrypted file
        key: Encryption key (generates new if None)
    
    Returns:
        Encryption key used (store securely!)
    """
    if key is None:
        key = Fernet.generate_key()
    
    cipher = Fernet(key)
    
    with open(input_path, 'rb') as f:
        plaintext = f.read()
    
    encrypted = cipher.encrypt(plaintext)
    
    with open(output_path, 'wb') as f:
        f.write(encrypted)
    
    return key

def decrypt_file(
    input_path: str,
    output_path: str,
    key: bytes
) -> None:
    """
    Decrypt file encrypted with Fernet.
    
    Args:
        input_path: Path to encrypted file
        output_path: Path to save decrypted file
        key: Encryption key
    
    Raises:
        cryptography.fernet.InvalidToken: If key is wrong
    """
    cipher = Fernet(key)
    
    with open(input_path, 'rb') as f:
        encrypted = f.read()
    
    decrypted = cipher.decrypt(encrypted)
    
    with open(output_path, 'wb') as f:
        f.write(decrypted)
''',

    "jwt_token": '''import jwt
from datetime import datetime, timedelta
from typing import Dict, Any, Optional

def generate_jwt_token(
    payload: Dict[str, Any],
    secret_key: str,
    expiry_hours: int = 24,
    algorithm: str = "HS256"
) -> str:
    """
    Generate JWT token with expiration.
    
    Args:
        payload: Data to encode in token
        secret_key: Secret key for signing
        expiry_hours: Token validity in hours
        algorithm: Signing algorithm
    
    Returns:
        Encoded JWT token string
    """
    expiry = datetime.utcnow() + timedelta(hours=expiry_hours)
    
    token_payload = {
        **payload,
        "exp": expiry,
        "iat": datetime.utcnow()
    }
    
    token = jwt.encode(token_payload, secret_key, algorithm=algorithm)
    return token

def verify_jwt_token(
    token: str,
    secret_key: str,
    algorithm: str = "HS256"
) -> Optional[Dict[str, Any]]:
    """
    Verify and decode JWT token.
    
    Args:
        token: JWT token string
        secret_key: Secret key for verification
        algorithm: Signing algorithm
    
    Returns:
        Decoded payload dict or None if invalid/expired
    """
    try:
        payload = jwt.decode(token, secret_key, algorithms=[algorithm])
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None
''',

    "redis_cache": '''import redis
import json
from typing import Any, Optional

def setup_redis_client(
    host: str = "localhost",
    port: int = 6379,
    db: int = 0,
    password: Optional[str] = None
) -> redis.Redis:
    """
    Create Redis client connection.
    
    Args:
        host: Redis server host
        port: Redis server port
        db: Database number
        password: Authentication password
    
    Returns:
        Redis client instance
    """
    return redis.Redis(
        host=host,
        port=port,
        db=db,
        password=password,
        decode_responses=True
    )

def cache_data(
    client: redis.Redis,
    key: str,
    value: Any,
    ttl_seconds: int = 3600
) -> bool:
    """
    Cache data in Redis with TTL.
    
    Args:
        client: Redis client instance
        key: Cache key
        value: Data to cache (will be JSON serialized)
        ttl_seconds: Time-to-live in seconds
    
    Returns:
        True if successful
    """
    serialized = json.dumps(value)
    return client.setex(key, ttl_seconds, serialized)

def get_cached_data(
    client: redis.Redis,
    key: str
) -> Optional[Any]:
    """
    Retrieve cached data from Redis.
    
    Args:
        client: Redis client instance
        key: Cache key
    
    Returns:
        Cached data or None if not found/expired
    """
    data = client.get(key)
    if data is None:
        return None
    return json.loads(data)
''',

    "s3_upload": '''import boto3
from typing import Optional
from botocore.exceptions import ClientError

def upload_file_to_s3(
    file_path: str,
    bucket_name: str,
    object_key: str,
    aws_access_key: Optional[str] = None,
    aws_secret_key: Optional[str] = None,
    region: str = "us-east-1"
) -> bool:
    """
    Upload file to AWS S3 bucket.
    
    Args:
        file_path: Local file path
        bucket_name: S3 bucket name
        object_key: S3 object key (path in bucket)
        aws_access_key: AWS access key ID (uses env if None)
        aws_secret_key: AWS secret access key (uses env if None)
        region: AWS region
    
    Returns:
        True if upload successful, False otherwise
    """
    try:
        if aws_access_key and aws_secret_key:
            s3_client = boto3.client(
                's3',
                aws_access_key_id=aws_access_key,
                aws_secret_access_key=aws_secret_key,
                region_name=region
            )
        else:
            s3_client = boto3.client('s3', region_name=region)
        
        s3_client.upload_file(file_path, bucket_name, object_key)
        return True
    except ClientError as e:
        print(f"Upload failed: {e}")
        return False

def download_file_from_s3(
    bucket_name: str,
    object_key: str,
    file_path: str,
    region: str = "us-east-1"
) -> bool:
    """
    Download file from AWS S3 bucket.
    
    Args:
        bucket_name: S3 bucket name
        object_key: S3 object key
        file_path: Local path to save file
        region: AWS region
    
    Returns:
        True if download successful, False otherwise
    """
    try:
        s3_client = boto3.client('s3', region_name=region)
        s3_client.download_file(bucket_name, object_key, file_path)
        return True
    except ClientError as e:
        print(f"Download failed: {e}")
        return False
''',
}


def get_template_for_prompt(prompt: str) -> Optional[str]:
    """
    根据提示词匹配最合适的模板
    
    Args:
        prompt: 用户提示词
    
    Returns:
        匹配的模板代码，如果没有匹配则返回 None
    """
    prompt_lower = prompt.lower()
    
    # SQLite 数据库操作
    if ("sqlite" in prompt_lower or "sql" in prompt_lower) and ("query" in prompt_lower or "database" in prompt_lower or "db" in prompt_lower):
        return TEMPLATES["sqlite_database"]
    
    # 文件加密/解密
    if ("encrypt" in prompt_lower or "decrypt" in prompt_lower or "cipher" in prompt_lower) and "file" in prompt_lower:
        return TEMPLATES["file_encryption"]
    
    # JWT token 生成/验证
    if "jwt" in prompt_lower or ("token" in prompt_lower and ("generate" in prompt_lower or "verify" in prompt_lower or "auth" in prompt_lower)):
        return TEMPLATES["jwt_token"]
    
    # Redis 缓存
    if "redis" in prompt_lower or ("cache" in prompt_lower and ("set" in prompt_lower or "get" in prompt_lower or "ttl" in prompt_lower)):
        return TEMPLATES["redis_cache"]
    
    # AWS S3 上传/下载
    if ("s3" in prompt_lower or "aws" in prompt_lower) and ("upload" in prompt_lower or "download" in prompt_lower or "bucket" in prompt_lower):
        return TEMPLATES["s3_upload"]
    
    # 日志轮转
    if "rotating" in prompt_lower and ("log" in prompt_lower or "file" in prompt_lower):
        return TEMPLATES["rotating_file_logger"]
    
    # HTTP 重试请求
    if ("http" in prompt_lower or "request" in prompt_lower or "url" in prompt_lower) and "retr" in prompt_lower:
        return TEMPLATES["http_retry_request"]
    
    # Pandas CSV 过滤
    if "pandas" in prompt_lower or ("csv" in prompt_lower and "filter" in prompt_lower):
        return TEMPLATES["pandas_filter_csv"]
    
    return None
